""" A collection of functions useful for analyzing structures of fullerenes. """
import sys        # call sys.exit() at errors.
import numpy as np

def get_ring(nm, nn_index, start_index, p):
    """ Returns rings of "nm" (5 or 6) members.

    The underlying algorithm is recursion. For most recursive algorithms,
    the execution speed is always a concern. But, in this program, the
    recusion will stop at the fifth (or sixth) atom counting from the
    starting atom, the number of searching actions is constant (2^5=32
    for a pentagon, 2^6=64 for a hexagon), and the total number of actions
    is thus proportional to the number of atoms N, being 32*N for pentagons,
    or 64*N for hexagons. The complexity is O(N), which should be efficient
    for large systems (i.e., large number of atoms).

    Arguments:
    nm: number of members. 
        This should be 5 for pentagons, or 6 for hexagons.
    nn_index: nearest-neighbor index.
        If there are N atoms, then the dimension of nn_index is N x (N-1).
    start_index: starting index for searching for rings.
        Usually, this function (get_ring()) is called iteratively, and 
        starting_index should iterate all atoms, one at a time.  
    p: array of rings.
        This array is gradually filled during the search.
        It maintains two arrays, one is unsorted, the other is sorted.
        The sorted array is for checking where a ring is already in the list.
    
    The present implementation entails recursion. I think it is inefficient
    because each atom is searched for many times that are unnecessary.
    The greatest difficulty is to identify repeating rings.
    """
    t = nn_index[start_index]   #t is an array storing indices of NN
    level = [1]  #Python cannot pass by reference, so I encapsulate it in a list
    stored_indices = []  #This array is updated dynamically using append and pop
    stored_indices.append(start_index)

    # Only call the first 3 NN, because in a fullerene molecule, each carbon
    # atom exactly has 3 first NN.
    # get_ring_() is called recursively. Other methods would be too verbose.
    if t[0] not in stored_indices:
        get_ring_(nm, nn_index, start_index, t[0], level, stored_indices, p)
    if t[1] not in stored_indices:
        get_ring_(nm, nn_index, start_index, t[1], level, stored_indices, p)
    if t[2] not in stored_indices:
        get_ring_(nm, nn_index, start_index, t[2], level, stored_indices, p)


def get_ring_(nm, nn_index, start_index, current_index, level, stored_indices, p):
    t = nn_index[current_index][0:3]
    level[0] = level[0]+1
    stored_indices.append(current_index)
    #print(current_index, t[0], t[1], t[2], level, stored_indices, p)

    if level[0]==nm:
        if t[0]==start_index or t[1]==start_index or t[2]==start_index :
            # found one
            ind = stored_indices[:]
            ind.sort()
            ind = tuple(ind)

            insertable = True
            for x in p[1]:
                if ind == x :
                    insertable = False
                    break
            if insertable : 
                p[0].append(tuple(stored_indices))
                p[1].append(tuple(ind))

            level[0] = level[0]-1
            stored_indices.pop()
            return
        else:
            # starting from here, the next search is sure to exceed the preset 
            # number of members or nm; stop here, and return.
            level[0] = level[0]-1
            stored_indices.pop()
            return
    else:
        # recursively call get_ring_() to continue the search.
        if t[0] not in stored_indices:
            get_ring_(nm, nn_index, start_index, t[0], level, stored_indices, p)
        if t[1] not in stored_indices:
            get_ring_(nm, nn_index, start_index, t[1], level, stored_indices, p)
        if t[2] not in stored_indices:
            get_ring_(nm, nn_index, start_index, t[2], level, stored_indices, p)
        
        level[0] = level[0]-1
        stored_indices.pop()
        return


def getPentagons(nn_index):
    """Returns all pentagons from the network described by nn_index."""
    p = [[], []]
    for i, nn in enumerate(nn_index):
        get_ring(5, nn_index, i, p)
    return p[0]

def getHexagons(nn_index):
    """Returns all hexagons from the network described by nn_index."""
    p = [[], []]
    for i, nn in enumerate(nn_index):
        get_ring(6, nn_index, i, p)
    return p[0]
    

def edge_ascending_indices(r):
    """ Returns edges of a ring, each edge is a pair, in ascending order."""
    ret = []
    for i in range(len(r)-1):
        if r[i]<r[i+1]: ret.append( (r[i], r[i+1]) )
        else : ret.append( (r[i+1], r[i]) )
    # the edges are cyclic; do not forget r[-1]--r[0]
    if r[-1]<r[0]: ret.append( (r[-1], r[0]) )
    else: ret.append( (r[0], r[-1]) )
    # now the edges ready
    return ret


def edgeIndices(pentagons, hexagons):
    """ Returns bond indicess of B66, B56, and B55 types."""
    B66_indices = []
    B56_indices = []
    B55_indices = []
    # search B66 bonds
    for i, p in enumerate(hexagons[:-1]) :
        b1 = edge_ascending_indices(p)
        for p2 in hexagons[i+1:]:
            b2 = edge_ascending_indices(p2)
            for b in b1 :
                if ( (b in b2) and (b not in B66_indices) ):
                    B66_indices.append(b)

    # search B56 bonds
    for p in pentagons :
        b1 = edge_ascending_indices(p)
        for p2 in hexagons:
            b2 = edge_ascending_indices(p2)
            for b in b1 :
                if ( (b in b2) and (b not in B56_indices) ):
                    B56_indices.append(b)
   
    # search B55 bonds
    for i, p in enumerate(pentagons[:-1]) :
        b1 = edge_ascending_indices(p)
        for p2 in pentagons[i+1:]:
            b2 = edge_ascending_indices(p2)
            for b in b1 :
                if ( (b in b2) and (b not in B66_indices) ):
                    B66_indices.append(b)

    return (B55_indices, B56_indices, B66_indices)
                


# A helper function for calculating coordinates of an adsorbed atom
def projecting_position(p1, p2, h):
    """ Returns a position along p1---p2 and is h away from p2.
    
    The line will be p1---p2---p_new, and the distance of p2---p_new
    is h. The coordinate of p_new is returned.
    The coordinates should be cartesian ones."""
    d = np.sqrt((p2[0]-p1[0])**2 + (p2[1]-p1[1])**2 + (p2[2]-p1[2])**2)
    t = h/d
    x = (t + 1) * p2[0] - t * p1[0]
    y = (t + 1) * p2[1] - t * p1[1]
    z = (t + 1) * p2[2] - t * p1[2]
    return (x, y, z)


# test code
if __name__=="__main__":
    from structure import VASP_structure
    s = VASP_structure()
    f_CONTCAR =  open("CONTCAR", "r")
    s.read(f_CONTCAR)
    f_CONTCAR.close()

    # Generates nearest-neighbor (NN) list.
    NN = s.NN()[0]
    pentagons = getPentagons(NN)
    hexagons = getHexagons(NN)
    # print("pentagons", pentagons)
    # print("hexagons", hexagons)
    
    # Calculates indice-pairs of B55, B56, and B66 bonds.
    B55, B56, B66 = edgeIndices(pentagons, hexagons)
    # print("B55", len(B55), B55)    
    # print("B56", len(B56), B56)
    # print("B66", len(B66), B66)

    # Calculates bond lengths of B56 and B66.
    # Most fullerene molecules do not have B55.
    dist_b56 = [s.distance(x[0], x[1]) for x in B56]
    dist_b66 = [s.distance(x[0], x[1]) for x in B66]
    # for d in dist_b56: print(d)
    # for d in dist_b66: print(d)

    # Calculates mean bond lengths.
    dist_b56_average = np.average(dist_b56)
    dist_b56_sd = np.std(dist_b56, ddof=1)
    dist_b66_average = np.average(dist_b66)
    dist_b66_sd = np.std(dist_b66, ddof=1)
    print("bond\taverage\tsd")
    print("b56\t{0:6.4f}\t{1:6.4f}\n".format(dist_b56_average, dist_b56_sd))
    print("b66\t{0:6.4f}\t{1:6.4f}\n".format(dist_b66_average, dist_b66_sd))

